#include "../../include/thread_pool.h"
#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include <unistd.h>

// 工作线程函数
static void *thread_worker(void *arg);

// 创建线程池
thread_pool_t *thread_pool_create(int min_threads, int max_threads, int max_queue)
{
    if (min_threads <= 0 || max_threads <= 0 || min_threads > max_threads || max_queue <= 0)
    {
        return NULL;
    }

    thread_pool_t *pool = (thread_pool_t *)malloc(sizeof(thread_pool_t));
    if (pool == NULL)
    {
        return NULL;
    }

    // 初始化线程池属性
    pool->thread_count = min_threads;
    pool->min_thread_count = min_threads;
    pool->max_thread_count = max_threads;
    pool->queue_size = 0;
    pool->max_queue_size = max_queue;
    pool->queue_head = NULL;
    pool->queue_tail = NULL;
    pool->shutdown = false;

    // 分配线程数组
    pool->threads = (pthread_t *)malloc(sizeof(pthread_t) * max_threads);
    if (pool->threads == NULL)
    {
        free(pool);
        return NULL;
    }

    // 初始化互斥锁和条件变量
    if (pthread_mutex_init(&pool->lock, NULL) != 0 ||
        pthread_cond_init(&pool->notify, NULL) != 0)
    {
        free(pool->threads);
        free(pool);
        return NULL;
    }

    // 创建工作线程
    for (int i = 0; i < min_threads; i++)
    {
        if (pthread_create(&pool->threads[i], NULL, thread_worker, (void *)pool) != 0)
        {
            // 创建线程失败，销毁线程池
            thread_pool_destroy(pool);
            return NULL;
        }
    }

    return pool;
}

// 添加任务到线程池
int thread_pool_add_task(thread_pool_t *pool, void (*function)(void *), void *argument)
{
    if (pool == NULL || function == NULL)
    {
        return -1;
    }

    // 获取互斥锁
    if (pthread_mutex_lock(&pool->lock) != 0)
    {
        return -1;
    }

    // 如果线程池已关闭或队列已满，返回错误
    if (pool->shutdown || pool->queue_size >= pool->max_queue_size)
    {
        pthread_mutex_unlock(&pool->lock);
        return -1;
    }

    // 创建新任务
    thread_task_t *task = (thread_task_t *)malloc(sizeof(thread_task_t));
    if (task == NULL)
    {
        pthread_mutex_unlock(&pool->lock);
        return -1;
    }

    task->function = function;
    task->argument = argument;
    task->next = NULL;

    // 将任务添加到队列尾部
    if (pool->queue_size == 0)
    {
        pool->queue_head = task;
        pool->queue_tail = task;
    }
    else
    {
        pool->queue_tail->next = task;
        pool->queue_tail = task;
    }

    pool->queue_size++;

    // 通知等待中的线程有新任务
    pthread_cond_signal(&pool->notify);
    pthread_mutex_unlock(&pool->lock);

    return 0;
}

// 销毁线程池
int thread_pool_destroy(thread_pool_t *pool)
{
    if (pool == NULL)
    {
        return -1;
    }

    // 获取互斥锁
    if (pthread_mutex_lock(&pool->lock) != 0)
    {
        return -1;
    }

    // 如果线程池已经关闭，直接返回
    if (pool->shutdown)
    {
        pthread_mutex_unlock(&pool->lock);
        return -1;
    }

    // 标记线程池为关闭状态
    pool->shutdown = true;

    // 释放互斥锁并唤醒所有等待的线程
    pthread_mutex_unlock(&pool->lock);
    pthread_cond_broadcast(&pool->notify);

    // 等待所有工作线程结束
    for (int i = 0; i < pool->thread_count; i++)
    {
        pthread_join(pool->threads[i], NULL);
    }

    // 清理线程池资源
    pthread_mutex_lock(&pool->lock);

    // 清理任务队列
    thread_task_t *task;
    while (pool->queue_head != NULL)
    {
        task = pool->queue_head;
        pool->queue_head = task->next;
        free(task);
    }

    pthread_mutex_unlock(&pool->lock);

    // 销毁互斥锁和条件变量
    pthread_mutex_destroy(&pool->lock);
    pthread_cond_destroy(&pool->notify);

    // 释放线程数组和线程池结构
    free(pool->threads);
    free(pool);

    return 0;
}

// 工作线程函数
static void *thread_worker(void *arg)
{
    thread_pool_t *pool = (thread_pool_t *)arg;
    thread_task_t *task;

    while (1)
    {
        // 获取互斥锁
        pthread_mutex_lock(&pool->lock);

        // 等待任务或关闭信号
        while (pool->queue_size == 0 && !pool->shutdown)
        {
            pthread_cond_wait(&pool->notify, &pool->lock);
        }

        // 如果线程池已关闭且队列为空，退出线程
        if (pool->shutdown && pool->queue_size == 0)
        {
            pthread_mutex_unlock(&pool->lock);
            pthread_exit(NULL);
        }

        // 获取队列头部的任务
        task = pool->queue_head;
        if (task != NULL)
        {
            pool->queue_head = task->next;
            if (pool->queue_head == NULL)
            {
                pool->queue_tail = NULL;
            }
            pool->queue_size--;
        }

        pthread_mutex_unlock(&pool->lock);

        // 如果有任务，执行它
        if (task != NULL)
        {
            (task->function)(task->argument);
            free(task);
        }
    }

    return NULL;
}